from fractions import Fraction
import itertools
from itertools import islice

from algorithm.list import binary_search
from number_theory.gcd import gcd
from number_theory.number_theory import even_factor
from number_theory.number_theory import jacobi
from number_theory.nth_root import int_sqrt


__all__ = ["is_prime", "is_probably_prime", "Prime"]



class Prime():
    small_primes = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37,
            41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97]

    @staticmethod
    def init_primes(bound):
        max_prime = Prime.small_primes[-1]
        if max_prime >= bound:
            return
        if (max_prime + 2) ** 2 < bound:
            Prime.init_primes(int_sqrt(bound))
        length = len(Prime.small_primes)
        for n in range(max_prime + 2, bound + 1, 2):
            if all(n % p for p in islice(Prime.small_primes, 1, length)):
                Prime.small_primes.append(n)

    @staticmethod
    def generate(bound = None):
        for p in Prime.small_primes:
            if bound and p > bound:
                return
            yield p
        for p in itertools.count(Prime.small_primes[-1] + 2, 2):
            if bound and p > bound:
                return
            elif is_prime(p):
                yield p


def is_prime(n):
    return is_probably_prime(n)


def is_probably_prime(n):
    if n <= 1:
        return False
    elif binary_search(Prime.small_primes, n):
        return True
    else:
        if any(n % p == 0 for p in Prime.small_primes):
            return False
        return test_miller_rabin(n) and test_lucas(n)


def test_miller_rabin(n, a_list = (2,)):
    for a in a_list:
        if gcd(a, n) != 1:
            return False

    m = n - 1
    s, d = even_factor(m)

    for a in a_list:
        x = pow(a, d, n)
        if x == 1 or x == m:
            continue
        for times in range(s - 1):
            x = (x ** 2) % n
            if x == 1:
                return False
            elif x == m:
                break
        else:
            return False
    return True


class _LucasMatrix:

    __slots__ = ("_11", "_12", "_21", "_22")

    def __init__(self, a, b, c, d):
        mod = type(self).mod
        self._11 = a % mod
        self._12 = b % mod
        self._21 = c % mod
        self._22 = d % mod

    @property
    def a11(self):
        return self._11

    def __imul__(self, other):
        mod = type(self).mod
        self._11, self._12, self._21, self._22 = (
                (self._11 * other._11 + self._12 * other._21) % mod,
                (self._11 * other._12 + self._12 * other._22) % mod,
                (self._21 * other._11 + self._22 * other._21) % mod,
                (self._21 * other._12 + self._22 * other._22) % mod)
        return self

    def __pow__(self, exp):
        base = self
        answer = type(self)(1, 0, 0, 1)
        while exp:
            if exp & 1:
                answer *= base
            base *= base
            exp >>= 1
        return answer


LucasMatrix = lambda mod: type(
        "LucasMatrix", (_LucasMatrix,), {"mod": mod})


def test_lucas(n):
    q = 1
    while True:
        q = -q if q > 0 else -q + 1
        if jacobi(1 - 4 * q, n) == -1:
            break
    return not (LucasMatrix(n)(1, -q, 1, 0) ** n).a11
